from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
import numpy as np
import matplotlib.pyplot as plt
import torch

plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号


def draw3d(X, vocab):
    # 创建一个新的图形
    fig = plt.figure(figsize=(16, 16))
    # 添加一个3D子图
    ax = fig.add_subplot(111, projection='3d')
    # 为每个点生成一个颜色（这里使用RGB颜色）
    colors = np.random.rand(X.shape[0], 3)  # 生成50个RGB颜色值，每个颜色值是一个长度为3的数组
    # 设置标题和轴标签
    ax.set_title('带有不同颜色和文字序号的3D散点图示例')
    # 绘制3D散点图
    x = X[:, 0] * 2
    y = X[:, 1] * 2
    z = X[:, 2] * 2

    scatter = ax.scatter(x, y, z, s=5, c=colors, marker='o')
    ax.set_xlabel('X轴')
    ax.set_ylabel('Y轴')
    ax.set_zlabel('Z轴')
    # 为每个数据点添加文字标签
    # vocab.to_tokens(0)
    for i, txt in enumerate(np.arange(1, X.shape[0] + 1)):  # 假设我们使用从1开始的序号作为标签
        ax.text(x[i], y[i], z[i], vocab[int(txt) - 1], fontsize=8, color="blue")  # 这里的txt是标签，可以根据需要自定义
    # 显示颜色条
    fig.colorbar(scatter)
    # 显示图形
    # plt.show()


def draw2d(X, vocab):
    # 创建一个新的图形
    fig = plt.figure()
    # 添加一个3D子图
    ax = fig.add_subplot(111)
    # 为每个点生成一个颜色（这里使用RGB颜色）
    colors = np.random.rand(X.shape[0], 1)  # 生成50个RGB颜色值，每个颜色值是一个长度为2的数组
    # 设置标题和轴标签
    ax.set_title('词向量可视化')
    # 绘制3D散点图
    x = X[:, 0]
    y = X[:, 1]

    scatter = ax.scatter(x, y, s=5, c=colors, marker='o')
    # 设置标题和轴标签
    ax.set_xlabel('X轴')
    ax.set_ylabel('Y轴')

    # 为每个数据点添加文字标签
    for i, txt in enumerate(np.arange(1, X.shape[0] + 1)):  # 假设我们使用从1开始的序号作为标签
        ax.text(x[i], y[i], vocab[int(txt) - 1], fontsize=8, color="blue")  # 这里的txt是标签，可以根据需要自定义
    # # 显示颜色条
    fig.colorbar(scatter)
    # 显示图形
    # plt.show()


def paca(X, dim=2):
    # 创建一个StandardScaler对象，用于对数据进行标准化
    scaler = StandardScaler()
    # 使用StandardScaler对象对数据进行标准化
    X_std = scaler.fit_transform(X)
    # 创建一个PCA对象，这里我们选择将数据降至一维
    pca = PCA(n_components=dim)
    # 使用PCA对数据进行降维
    X_pca = pca.fit_transform(X_std)
    # 返回降维后的数据
    return X_pca


# vocab.to_tokens([100])
def draw_vocab_pca(X, vocab, dim=2):
    if dim == 2:
        text_list = [vocab.to_tokens([i])[0] for i in range(len(X))]
        draw2d(X, text_list)
    else:
        text_list = [vocab.to_tokens([i])[0] for i in range(len(X))]
        draw3d(X, text_list)


def draw_words_pca(p_words, X, vocab, dim=2):
    new_X = []
    text_list = []
    X = paca(X, 3)
    # print(X.shape)
    X = [torch.tensor(i) for i in X]

    for i in p_words:
        for j in i:
            new_X.append(X[vocab[j]])
            text_list.append(j)

    new_X = torch.stack(new_X, dim=0)
    if dim == 2:
        draw2d(new_X, text_list)
    else:
        draw3d(new_X, text_list)

# X = torch.randn(
#     6, 200)
# print(paca(X, dim=3))
